﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace FunasrSTT.Translate.Services;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;

/// <summary>
/// 多任务管理器
/// </summary>
public class TaskManager : IDisposable
{
    private readonly ConcurrentDictionary<Guid, TaskController> _tasks = new();
    private readonly CancellationTokenSource _globalCancellation = new();
    private bool _disposed;
    private int _maxConcurrentTasks = Environment.ProcessorCount * 2;
    private readonly SemaphoreSlim _concurrencyLimiter;

    public event EventHandler<TaskStatusChangedEventArgs> TaskStatusChanged;
    public event EventHandler<TaskExceptionEventArgs> TaskErrorOccurred;

    /// <summary>
    /// 初始化任务管理器
    /// </summary>
    /// <param name="maxConcurrentTasks"></param>
    public TaskManager(int maxConcurrentTasks = -1)
    {
        if (maxConcurrentTasks > 0) _maxConcurrentTasks = maxConcurrentTasks;
        _concurrencyLimiter = new SemaphoreSlim(_maxConcurrentTasks, _maxConcurrentTasks);
    }

    public Guid Start(
        Func<CancellationToken, Task> asyncAction,
        TaskPriority priority = TaskPriority.Normal,
        string taskName = null,
        TaskRestartPolicy restartPolicy = TaskRestartPolicy.Never,
        TimeSpan? timeout = null)
    {
        ThrowIfDisposed();

        var taskId = Guid.NewGuid();
        var cts = CancellationTokenSource.CreateLinkedTokenSource(_globalCancellation.Token);

        var taskInfo = new TaskController(
            taskId: taskId,
            name: taskName ?? $"Task_{taskId.ToString()[..8]}",
            priority: priority,
            restartPolicy: restartPolicy,
            cancellationSource: cts,
            timeout: timeout
        );

        var task = Task.Run(async () =>
        {
            await _concurrencyLimiter.WaitAsync(cts.Token);
            try
            {
                await ExecuteTaskAsync(asyncAction, taskInfo, cts);
            }
            finally
            {
                _concurrencyLimiter.Release();
            }
        }, cts.Token);

        taskInfo.Task = task;
        _tasks.TryAdd(taskId, taskInfo);
        OnTaskStatusChanged(taskInfo, TaskStatus.Running);

        return taskId;
    }

    private async Task ExecuteTaskAsync(
        Func<CancellationToken, Task> asyncAction,
        TaskController taskInfo,
        CancellationTokenSource cts)
    {
        int attempt = 0;
        while (!cts.IsCancellationRequested)
        {
            try
            {
                attempt++;
                taskInfo.AttemptCount = attempt;
                OnTaskStatusChanged(taskInfo, TaskStatus.Running);

                using var timeoutCts = taskInfo.Timeout.HasValue
                    ? CancellationTokenSource.CreateLinkedTokenSource(cts.Token)
                    : null;

                if (timeoutCts != null)
                {
                    timeoutCts.CancelAfter(taskInfo.Timeout.Value);
                }

                var token = timeoutCts?.Token ?? cts.Token;

                await asyncAction(token);

                taskInfo.LastStatus = TaskStatus.Completed;
                OnTaskStatusChanged(taskInfo, TaskStatus.Completed);
                return; // Success, exit task
            }
            catch (OperationCanceledException)
            {
                taskInfo.LastStatus = TaskStatus.Canceled;
                OnTaskStatusChanged(taskInfo, TaskStatus.Canceled);
                throw;
            }
            catch (Exception ex)
            {
                taskInfo.LastStatus = TaskStatus.Faulted;
                taskInfo.LastException = ex;
                OnTaskErrorOccurred(taskInfo, ex);
                OnTaskStatusChanged(taskInfo, TaskStatus.Faulted);

                if (!ShouldRetry(taskInfo.RestartPolicy, attempt))
                {
                    return; // No more retries
                }
            }

            // Wait before retrying
            if (taskInfo.RestartPolicy == TaskRestartPolicy.WithDelay)
            {
                await Task.Delay(CalculateBackoff(attempt), cts.Token);
            }
        }
    }

    private bool ShouldRetry(TaskRestartPolicy policy, int attempt)
    {
        return policy switch
        {
            TaskRestartPolicy.Never => false,
            TaskRestartPolicy.Always => true,
            TaskRestartPolicy.OnFailure => attempt < 3,
            TaskRestartPolicy.WithDelay => attempt < 5,
            _ => false
        };
    }

    private TimeSpan CalculateBackoff(int attempt)
    {
        return TimeSpan.FromSeconds(Math.Pow(2, Math.Min(attempt, 5)));
    }

    public void Stop(Guid taskId)
    {
        ThrowIfDisposed();

        if (_tasks.TryGetValue(taskId, out var controller))
        {
            controller.CancellationSource.Cancel();
        }
    }

    public async Task StopAsync(Guid taskId, TimeSpan? timeout = null)
    {
        ThrowIfDisposed();

        if (!_tasks.TryGetValue(taskId, out var controller)) return;

        controller.CancellationSource.Cancel();

        try
        {
            var task = controller.Task;
            if (task != null && !task.IsCompleted)
            {
                await (timeout.HasValue
                    ? task.WaitAsync(timeout.Value)
                    : task);
            }
        }
        catch (OperationCanceledException)
        {
            // Expected
        }
        finally
        {
            _tasks.TryRemove(taskId, out _);
        }
    }

    public async Task StopAllAsync(TimeSpan? timeout = null)
    {
        ThrowIfDisposed();

        // Cancel all tasks
        _globalCancellation.Cancel();

        try
        {
            // Wait for all tasks to complete
            await Task.WhenAll(GetActiveTasks()).WaitAsync(
                timeout ?? TimeSpan.FromSeconds(15));
        }
        catch (TimeoutException)
        {
            // Log or handle timeout
        }
        finally
        {
            // Reset cancellation token
            _globalCancellation.Dispose();
            //_globalCancellation = new CancellationTokenSource();
            _tasks.Clear();
        }
    }

    public void Pause(Guid taskId)
    {
        ThrowIfDisposed();

        if (_tasks.TryGetValue(taskId, out var controller))
        {
            controller.IsPaused = true;
            OnTaskStatusChanged(controller, TaskStatus.Paused);
        }
    }

    public void Resume(Guid taskId)
    {
        ThrowIfDisposed();

        if (_tasks.TryGetValue(taskId, out var controller))
        {
            controller.IsPaused = false;
            OnTaskStatusChanged(controller, TaskStatus.Running);
        }
    }

    public TaskInfo GetTaskInfo(Guid taskId)
    {
        ThrowIfDisposed();

        return _tasks.TryGetValue(taskId, out var controller)
            ? controller.GetInfo()
            : null;
    }

    public IEnumerable<TaskInfo> GetAllTasksInfo()
    {
        ThrowIfDisposed();

        foreach (var controller in _tasks.Values)
        {
            yield return controller.GetInfo();
        }
    }

    public int GetActiveTaskCount()
    {
        ThrowIfDisposed();

        int count = 0;
        foreach (var controller in _tasks.Values)
        {
            if (controller.Status == TaskStatus.Running) count++;
        }
        return count;
    }

    private IEnumerable<Task> GetActiveTasks()
    {
        foreach (var controller in _tasks.Values)
        {
            if (controller.Task != null && !controller.Task.IsCompleted)
            {
                yield return controller.Task;
            }
        }
    }

    private void OnTaskStatusChanged(TaskController taskInfo, TaskStatus status)
    {
        taskInfo.LastStatus = status;
        TaskStatusChanged?.Invoke(this, new TaskStatusChangedEventArgs(
            taskInfo.Id,
            taskInfo.Name,
            status
        ));
    }

    private void OnTaskErrorOccurred(TaskController taskInfo, Exception exception)
    {
        TaskErrorOccurred?.Invoke(this, new TaskExceptionEventArgs(
            taskInfo.Id,
            taskInfo.Name,
            exception
        ));
    }

    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    protected virtual void Dispose(bool disposing)
    {
        if (_disposed) return;

        if (disposing)
        {
            _globalCancellation.Dispose();
            _concurrencyLimiter.Dispose();
            _globalCancellation.Cancel();
            Task.WhenAll(GetActiveTasks()).Wait(TimeSpan.FromSeconds(5));
        }

        _disposed = true;
    }

    private void ThrowIfDisposed()
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(TaskManager));
        }
    }

    ~TaskManager()
    {
        Dispose(false);
    }

    private class TaskController
    {
        public Guid Id { get; }
        public string Name { get; }
        public TaskPriority Priority { get; }
        public TaskRestartPolicy RestartPolicy { get; }
        public CancellationTokenSource CancellationSource { get; }
        public TimeSpan? Timeout { get; }
        public Task Task { get; set; }
        public int AttemptCount { get; set; }
        public DateTime StartTime { get; } = DateTime.UtcNow;
        public DateTime? EndTime { get; private set; }
        public TaskStatus LastStatus { get; set; } = TaskStatus.Created;
        public Exception LastException { get; set; }
        public bool IsPaused { get; set; }

        public TaskStatus Status
        {
            get
            {
                if (IsPaused) return TaskStatus.Paused;
                if (Task == null) return TaskStatus.Created;
                return Task.IsCompleted ? LastStatus : TaskStatus.Running;
            }
        }

        public TaskController(
            Guid taskId,
            string name,
            TaskPriority priority,
            TaskRestartPolicy restartPolicy,
            CancellationTokenSource cancellationSource,
            TimeSpan? timeout)
        {
            Id = taskId;
            Name = name;
            Priority = priority;
            RestartPolicy = restartPolicy;
            CancellationSource = cancellationSource;
            Timeout = timeout;
        }

        public TaskInfo GetInfo()
        {
            return new TaskInfo(
                Id,
                Name,
                Status,
                StartTime,
                EndTime,
                EndTime.HasValue ? EndTime.Value - StartTime : TimeSpan.Zero,
                AttemptCount,
                LastException
            );
        }
    }
}

public class TaskInfo
{
    public Guid Id { get; }
    public string Name { get; }
    public TaskStatus Status { get; }
    public DateTime StartTime { get; }
    public DateTime? EndTime { get; }
    public TimeSpan ExecutionTime { get; }
    public int AttemptCount { get; }
    public Exception LastException { get; }

    public TaskInfo(
        Guid id,
        string name,
        TaskStatus status,
        DateTime startTime,
        DateTime? endTime,
        TimeSpan executionTime,
        int attemptCount,
        Exception lastException)
    {
        Id = id;
        Name = name;
        Status = status;
        StartTime = startTime;
        EndTime = endTime;
        ExecutionTime = executionTime;
        AttemptCount = attemptCount;
        LastException = lastException;
    }
}

public class TaskStatusChangedEventArgs : EventArgs
{
    public Guid TaskId { get; }
    public string TaskName { get; }
    public TaskStatus Status { get; }

    public TaskStatusChangedEventArgs(Guid taskId, string taskName, TaskStatus status)
    {
        TaskId = taskId;
        TaskName = taskName;
        Status = status;
    }
}

public class TaskExceptionEventArgs : EventArgs
{
    public Guid TaskId { get; }
    public string TaskName { get; }
    public Exception Exception { get; }

    public TaskExceptionEventArgs(Guid taskId, string taskName, Exception exception)
    {
        TaskId = taskId;
        TaskName = taskName;
        Exception = exception;
    }
}

public enum TaskStatus
{
    Created,
    Running,
    Completed,
    Faulted,
    Canceled,
    Paused
}

public enum TaskPriority
{
    Low,
    Normal,
    High,
    Critical
}

public enum TaskRestartPolicy
{
    Never,
    Always,
    OnFailure,
    WithDelay
}

public static class TaskExtensions
{
    public static async Task WaitAsync(this Task task, TimeSpan timeout)
    {
        var delayTask = Task.Delay(timeout);
        var completedTask = await Task.WhenAny(task, delayTask);
        if (completedTask == delayTask)
        {
            throw new TimeoutException("The operation has timed out.");
        }

        await task; // Ensure any exceptions are thrown
    }
}